#include <assert.h>
#include <elf.h>
#include <errno.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define IMAGE_FILE "./image"
#define ARGS "[--extended] [--vm] <bootblock> <executable-file> ..."

#define OS_SIZE_LOC 0x1fc
#define SECTOR_SIZE 512
#define MAXSIZE_KERNEL 0x200000

/* structure to store command line options */
static struct {
    int vm;
    int extended;
} options;

/* prototypes of local functions */
static void create_image(int nfiles, char *files[], int extended);
static void error(char *fmt, ...);
static void read_ehdr(Elf64_Ehdr * ehdr, FILE * fp);
static void read_phdr(Elf64_Phdr * phdr, FILE * fp, int ph,
                      Elf64_Ehdr ehdr);
static void write_segment(Elf64_Ehdr ehdr, Elf64_Phdr phdr, FILE * fp,
                          FILE * img, int *nbytes, int *first, int extended);
static void write_os_size(int nbytes, FILE * img, int extended);

int main(int argc, char **argv)
{
    char *progname = argv[0];

    /* process command line options */
    options.vm = 0;
    options.extended = 0;
    while ((argc > 1) && (argv[1][0] == '-') && (argv[1][1] == '-')) {
        char *option = &argv[1][2];

        if (strcmp(option, "vm") == 0) {
            options.vm = 1;
        } else if (strcmp(option, "extended") == 0) {
            options.extended = 1;
        } else {
            error("%s: invalid option\nusage: %s %s\n", progname,
                  progname, ARGS);
        }
        argc--;
        argv++;
    }
    if (options.vm == 1) {
        error("%s: option --vm not implemented\n", progname);
    }
    if (argc < 3) {
        /* at least 3 args (createimage bootblock kernel) */
        error("usage: %s %s\n", progname, ARGS);
    }
    create_image(argc - 1, argv + 1, options.extended);
    return 0;
}

static void create_image(int nfiles, char *files[], int extended)
{
    int ph, nbytes = 0, first = 0;
    FILE *fp, *img;
    Elf64_Ehdr ehdr;
    Elf64_Phdr phdr;

    /* open the image file */
    img = fopen(IMAGE_FILE, "w+");

    /* for each input file */
    while (nfiles-- > 0) {

        /* open input file */
        fp = fopen(*files, "r");

        /* read ELF header */
        read_ehdr(&ehdr, fp);
        printf("0x%04lx: %s\n", ehdr.e_entry, *files);

        /* for each program header */
        for (ph = 0; ph < ehdr.e_phnum; ph++) {

            /* read program header */
            read_phdr(&phdr, fp, ph, ehdr);

            /*print extended infomation of segment*/
            if (extended)
            {
                printf("\tsegment %d\n",ph);
                printf("\t\toffset 0x%04lx\t\tvaddr 0x%04lx\n", phdr.p_offset, phdr.p_vaddr);
                printf("\t\tfilesz 0x%04lx\t\tmemsz 0x%04lx\n", phdr.p_filesz, phdr.p_memsz);
            }
            
            /* write segment to the image */
            write_segment(ehdr, phdr, fp, img, &nbytes, &first, extended);
        }
        fclose(fp);
        files++;
    }
    write_os_size(nbytes, img, extended);
    fclose(img);
}

static void read_ehdr(Elf64_Ehdr * ehdr, FILE * fp)
{
    fread(ehdr, sizeof(Elf64_Ehdr), 1, fp);
}

static void read_phdr(Elf64_Phdr * phdr, FILE * fp, int ph,
                      Elf64_Ehdr ehdr)
{
    long offset = ehdr.e_phoff + ehdr.e_phentsize * ph;
    fseek(fp, offset, SEEK_SET);
    fread(phdr, sizeof(Elf64_Phdr), 1, fp);
}

static void write_segment(Elf64_Ehdr ehdr, Elf64_Phdr phdr, FILE * fp,
                          FILE * img, int *nbytes, int *first, int extended)
{
    char temp_buf[MAXSIZE_KERNEL];

    long offset_r = phdr.p_offset;
    size_t memsz = phdr.p_memsz;
    fseek(fp, offset_r, SEEK_SET);
    fread(temp_buf, memsz, 1, fp);

    long offset_w = *first * SECTOR_SIZE;
    fseek(img, offset_w, SEEK_SET);
    fwrite(temp_buf, memsz, 1, img);

    int sector_count;   //count from 1
    int pad_num;       //number of padding char
    if (*first == 0)
    {
        (*first)++;
        sector_count = 1;
        pad_num = sector_count * SECTOR_SIZE - memsz;
    }
    else
    {
        *nbytes += memsz;
        sector_count = 1 + memsz / SECTOR_SIZE + (memsz % SECTOR_SIZE ? 1 : 0);
        pad_num = (sector_count - 1) * SECTOR_SIZE - memsz; 
    }
    /*padding*/
    char pad_buf[SECTOR_SIZE];
    for (int i = 0; i < SECTOR_SIZE; i++)
    {
        pad_buf[i] = 0;
    }
    fseek(img, offset_w + memsz, SEEK_SET);
    fwrite(pad_buf, 1, pad_num, img);

    /*extended info*/
    if (extended && memsz)
    {
        printf("\t\twriting 0x%04lx bytes\n", memsz);
        printf("\t\tpadding up to 0x%04lx\n", (long unsigned)sector_count * SECTOR_SIZE);        
    }
}

static void write_os_size(int nbytes, FILE * img, int extended)
{   
    int sector_count = nbytes / SECTOR_SIZE + (nbytes % SECTOR_SIZE ? 1 : 0);
    if (extended)
    {
        printf("os_size: %d sector(s)\n",sector_count);
    }
    
    char temp_buf[2];
    temp_buf[0] = (char)sector_count;
    temp_buf[1] = (char)(sector_count >> sizeof(char));

    fseek(img, OS_SIZE_LOC, SEEK_SET);
    fwrite(temp_buf, 2, 1, img);
}

/* print an error message and exit */
static void error(char *fmt, ...)
{
    va_list args;

    va_start(args, fmt);
    vfprintf(stderr, fmt, args);
    va_end(args);
    if (errno != 0) {
        perror(NULL);
    }
    exit(EXIT_FAILURE);
}
